from flask import Flask, jsonify
import multiprocessing as mp
from main import start_train
import traceback
from utils.log_output import simple_text_log

# Define global variables
app = Flask(__name__)
global_stats = {}

def worker(status):
    print('main: Worker start!')

    try:
        # Start train
        start_train()
        simple_text_log('train', 'The GPU server has completed training and updated the database.')
    except Exception as err:
        simple_text_log('error', str(err))
        s = traceback.format_exc()
        simple_text_log('error', s)
        simple_text_log('train', 'Training terminated due to some errors.')

    print('main: Worker finished.')
    status.value = 0

@app.route('/check-server', methods=['POST', 'GET'])
def check_server():
    # Check status
    if 'run_status' in global_stats:
        if global_stats['run_status'].value != 0:
            return jsonify({'status': 'running'})
    return jsonify({'status': 'normal'})

@app.route('/start-server', methods=['POST', 'GET'])
def start_server():
    # Check status
    if 'run_status' in global_stats:
        print(f'run_status is {global_stats["run_status"].value}')
        if global_stats['run_status'].value != 0:
            return jsonify({'success': 0, 'msg': 'Still training. Please stop first.'})
        cur_status = global_stats["run_status"]
        cur_status.value = 1
    else:
        print('cur_status not found')
        cur_status = mp.Manager().Value('cur_status', 1)
        global_stats['run_status'] = cur_status
    
    process = mp.Process(target=worker, args=(cur_status, ))
    process.start()
    global_stats['process'] = process
    return jsonify({'success': 1, 'msg': 'Training started succesfully!'})

@app.route('/stop-server', methods=['POST', 'GET'])
def stop_server():
    process: mp.Process = global_stats.get('process', None)
    if process is not None:
        process.kill()
        global_stats['run_status'].value = 0
        simple_text_log('train', 'Training stopped by user.')
    return jsonify({'success': 1, 'msg': 'Training stopped.'})

if __name__ == '__main__':
    simple_text_log('train', 'GPU Server started.')
    app.run('0.0.0.0', port=6789, debug=False)
